import asyncio
from typing import Optional
from mcp import ClientSession
from mcp.client.sse import sse_client


class MCPClient:

    def __init__(self):
        self.session: Optional[ClientSession] = None
        self._streams_context = None
        self._session_context = None

    async def connect_to_sse_server(self, server_url: str):
        """
        Connect to an MCP server running with SSE transport
        """
        self._streams_context = sse_client(url=server_url)
        streams = await self._streams_context.__aenter__()
        self._session_context = ClientSession(*streams)
        self.session: ClientSession = await self._session_context.__aenter__()
        await self.session.initialize()

    async def list_tools(self):
        """mcp工具列表"""
        response = await self.session.list_tools()
        tools = [tool.name for tool in response.tools]
        return tools

    async def cleanup(self):
        """Properly clean up the session and streams"""
        if self._session_context:
            await self._session_context.__aexit__(None, None, None)
        if self._streams_context:
            await self._streams_context.__aexit__(None, None, None)

    async def call_tool(self, tool_name, tool_args):
        """调用工具"""
        result = await self.session.call_tool(tool_name, tool_args)
        json_data = result.model_dump()
        return json_data

sse_url = 'http://127.0.0.1:8000/sse'

async def list_mcp_tools_test():
    print("======== 测试: 获取工具列表 ==========")
    client = MCPClient()
    await client.connect_to_sse_server(sse_url)
    tools = await client.list_tools()
    print("工具列表: ", tools)
    await client.cleanup()

async def sequential_thinking_test():
    print("======== 测试: 序列化思考 ==========")
    client = MCPClient()
    await client.connect_to_sse_server(sse_url)

    result = await client.call_tool(tool_name='sequentialthinking',
                                    tool_args={
                                        "thought": "将`泡一杯茶`分拆为多个可执行的步骤",
                                        "thoughtNumber": 1,
                                        "totalThoughts": 5,
                                        "nextThoughtNeeded": True,
                                    })
    for item in result["content"]:
        print(item['text'])
    await client.cleanup()


if __name__ == "__main__":
    asyncio.run(list_mcp_tools_test())          # 获取mcp工具列表
    asyncio.run(sequential_thinking_test())     # 序列化思考
